#pragma once

#include <numeric>
#include "det/clipper/clipper.hpp"

// locality-aware NMS
namespace lanms {

namespace cl = ClipperLib;

struct Polygon {
  cl::Path poly;
  float score;
};

float paths_area(const ClipperLib::Paths &ps) {
  float area = 0;
  for (auto &&p : ps) area += cl::Area(p);
  return area;
}

float poly_iou(const Polygon &a, const Polygon &b) {
  cl::Clipper clpr;
  clpr.AddPath(a.poly, cl::ptSubject, true);
  clpr.AddPath(b.poly, cl::ptClip, true);

  cl::Paths inter, uni;
  clpr.Execute(cl::ctIntersection, inter, cl::pftEvenOdd);
  clpr.Execute(cl::ctUnion, uni, cl::pftEvenOdd);

  auto inter_area = paths_area(inter), uni_area = paths_area(uni);
  return std::abs(inter_area) / std::max(std::abs(uni_area), 1.0f);
}

bool should_merge(const Polygon &a, const Polygon &b, float iou_threshold) {
  return poly_iou(a, b) > iou_threshold;
}

/**
 * Incrementally merge polygons
 */
class PolyMerger {
 public:
  PolyMerger() : score(0), nr_polys(0) { memset(data, 0, sizeof(data)); }

  /**
   * Add a new polygon to be merged.
   */
  void add(const Polygon &p_given) {
    Polygon p;
    if (nr_polys > 0) {
      // vertices of two polygons to merge may not in the same order;
      // we match their vertices by choosing the ordering that
      // minimizes the total squared distance.
      // see function normalize_poly for details.
      p = normalize_poly(get(), p_given);
    } else {
      p = p_given;
    }
    assert(p.poly.size() == 4);
    auto &poly = p.poly;
    auto s = p.score;
    data[0] += poly[0].X * s;
    data[1] += poly[0].Y * s;

    data[2] += poly[1].X * s;
    data[3] += poly[1].Y * s;

    data[4] += poly[2].X * s;
    data[5] += poly[2].Y * s;

    data[6] += poly[3].X * s;
    data[7] += poly[3].Y * s;

    score += p.score;

    nr_polys += 1;
  }

  inline std::int64_t sqr(std::int64_t x) { return x * x; }

  Polygon normalize_poly(const Polygon &ref, const Polygon &p) {
    std::int64_t min_d = std::numeric_limits<std::int64_t>::max();
    size_t best_start = 0, best_order = 0;

    for (size_t start = 0; start < 4; start++) {
      size_t j = start;
      std::int64_t d = (sqr(ref.poly[(j + 0) % 4].X - p.poly[(j + 0) % 4].X) +
                        sqr(ref.poly[(j + 0) % 4].Y - p.poly[(j + 0) % 4].Y) +
                        sqr(ref.poly[(j + 1) % 4].X - p.poly[(j + 1) % 4].X) +
                        sqr(ref.poly[(j + 1) % 4].Y - p.poly[(j + 1) % 4].Y) +
                        sqr(ref.poly[(j + 2) % 4].X - p.poly[(j + 2) % 4].X) +
                        sqr(ref.poly[(j + 2) % 4].Y - p.poly[(j + 2) % 4].Y) +
                        sqr(ref.poly[(j + 3) % 4].X - p.poly[(j + 3) % 4].X) +
                        sqr(ref.poly[(j + 3) % 4].Y - p.poly[(j + 3) % 4].Y));
      if (d < min_d) {
        min_d = d;
        best_start = start;
        best_order = 0;
      }

      d = (sqr(ref.poly[(j + 0) % 4].X - p.poly[(j + 3) % 4].X) +
           sqr(ref.poly[(j + 0) % 4].Y - p.poly[(j + 3) % 4].Y) +
           sqr(ref.poly[(j + 1) % 4].X - p.poly[(j + 2) % 4].X) +
           sqr(ref.poly[(j + 1) % 4].Y - p.poly[(j + 2) % 4].Y) +
           sqr(ref.poly[(j + 2) % 4].X - p.poly[(j + 1) % 4].X) +
           sqr(ref.poly[(j + 2) % 4].Y - p.poly[(j + 1) % 4].Y) +
           sqr(ref.poly[(j + 3) % 4].X - p.poly[(j + 0) % 4].X) +
           sqr(ref.poly[(j + 3) % 4].Y - p.poly[(j + 0) % 4].Y));
      if (d < min_d) {
        min_d = d;
        best_start = start;
        best_order = 1;
      }
    }

    Polygon r;
    r.poly.resize(4);
    auto j = best_start;
    if (best_order == 0) {
      for (size_t i = 0; i < 4; i++) r.poly[i] = p.poly[(j + i) % 4];
    } else {
      for (size_t i = 0; i < 4; i++) r.poly[i] = p.poly[(j + 4 - i - 1) % 4];
    }
    r.score = p.score;
    return r;
  }

  Polygon get() const {
    Polygon p;

    auto &poly = p.poly;
    poly.resize(4);
    auto score_inv = 1.0f / std::max(1e-8f, score);
    poly[0].X = data[0] * score_inv;
    poly[0].Y = data[1] * score_inv;
    poly[1].X = data[2] * score_inv;
    poly[1].Y = data[3] * score_inv;
    poly[2].X = data[4] * score_inv;
    poly[2].Y = data[5] * score_inv;
    poly[3].X = data[6] * score_inv;
    poly[3].Y = data[7] * score_inv;

    assert(score > 0);
    p.score = score;

    return p;
  }

 private:
  std::int64_t data[8];
  float score;
  std::int32_t nr_polys;
};

/**
 * The standard NMS algorithm.
 */
std::vector<Polygon> standard_nms(std::vector<Polygon> &polys,
                                  float iou_threshold) {
  size_t n = polys.size();
  if (n == 0) return {};
  std::vector<size_t> indices(n);
  std::iota(std::begin(indices), std::end(indices), 0);
  std::sort(std::begin(indices), std::end(indices), [&](size_t i, size_t j) {
    return polys[i].score > polys[j].score;
  });

  std::vector<size_t> keep;
  while (indices.size()) {
    size_t p = 0, cur = indices[0];
    keep.emplace_back(cur);
    for (size_t i = 1; i < indices.size(); i++) {
      if (!should_merge(polys[cur], polys[indices[i]], iou_threshold)) {
        indices[p++] = indices[i];
      }
    }
    indices.resize(p);
  }

  std::vector<Polygon> ret;
  for (auto &&i : keep) {
    ret.emplace_back(polys[i]);
  }
  return ret;
}

std::vector<Polygon> merge_quadrangle_n9(const float *data, size_t n,
                                         float iou_threshold) {
  using cInt = cl::cInt;

  // first pass
  std::vector<Polygon> polys;
  for (size_t i = 0; i < n; i++) {
    auto p = data + i * 9;
    Polygon poly{
        {
            {cInt(p[0]), cInt(p[1])},
            {cInt(p[2]), cInt(p[3])},
            {cInt(p[4]), cInt(p[5])},
            {cInt(p[6]), cInt(p[7])},
        },
        p[8],
    };

    if (polys.size()) {
      // merge with the last one
      auto &bpoly = polys.back();
      if (should_merge(poly, bpoly, iou_threshold)) {
        PolyMerger merger;
        merger.add(bpoly);
        merger.add(poly);
        bpoly = merger.get();
      } else {
        polys.emplace_back(poly);
      }
    } else {
      polys.emplace_back(poly);
    }
  }
  return standard_nms(polys, iou_threshold);
}

std::vector<std::vector<float>> polys2floats(
    const std::vector<lanms::Polygon> &polys) {
  std::vector<std::vector<float>> ret;
  for (size_t i = 0; i < polys.size(); i++) {
    auto &p = polys[i];
    auto &poly = p.poly;
    ret.emplace_back(std::vector<float>{
        float(poly[0].X),
        float(poly[0].Y),
        float(poly[1].X),
        float(poly[1].Y),
        float(poly[2].X),
        float(poly[2].Y),
        float(poly[3].X),
        float(poly[3].Y),
        float(p.score),
    });
  }

  return ret;
}

std::vector<std::vector<float>> polys2floats_new(
    std::vector<lanms::Polygon> &polys) {
  std::vector<std::vector<float>> ret;
  for (size_t i = 0; i < polys.size(); i++) {
    auto &p = polys[i];
    auto &poly = p.poly;
    ret.emplace_back(std::vector<float>{
        float(static_cast<double>(poly[0].X) / 10000.0),
        float(static_cast<double>(poly[0].Y) / 10000.0),
        float(static_cast<double>(poly[1].X) / 10000.0),
        float(static_cast<double>(poly[1].Y) / 10000.0),
        float(static_cast<double>(poly[2].X) / 10000.0),
        float(static_cast<double>(poly[2].Y) / 10000.0),
        float(static_cast<double>(poly[3].X) / 10000.0),
        float(static_cast<double>(poly[3].Y) / 10000.0),
        float(p.score),
    });
  }
  return ret;
}

}  // namespace lanms
